## 根据mask文件名获取图片和标签文件
import os
import shutil

def process_files(mask_folder, source_folder, target_base_folder):
    """
    根据mask文件名处理对应的标注文件和图片文件
    
    参数:
    mask_folder: 存放mask文件的文件夹路径
    source_folder: 包含Annotations和JPEGImages的源文件夹路径
    target_base_folder: 目标文件夹的根目录路径
    """
    # 创建目标文件夹
    target_anno_folder = os.path.join(target_base_folder, 'Annotations')
    target_jpeg_folder = os.path.join(target_base_folder, 'JPEGImages')
    
    # 如果目标文件夹不存在，创建它们
    os.makedirs(target_anno_folder, exist_ok=True)
    os.makedirs(target_jpeg_folder, exist_ok=True)
    
    # 源文件夹路径
    source_anno_folder = os.path.join(source_folder, 'Annotations')
    source_jpeg_folder = os.path.join(source_folder, 'JPEGImages')
    
    # 遍历mask文件夹
    for mask_file in os.listdir(mask_folder):
        if mask_file.startswith('mask_') and mask_file.endswith('.png'):
            # 提取中间的文件名部分 (例如: frame_01991)
            base_name = mask_file[5:-4]  # 去掉'mask_'前缀和'.png'后缀
            
            # 构建源文件路径
            source_anno_file = os.path.join(source_anno_folder, f"{base_name}.xml")
            source_jpeg_file = os.path.join(source_jpeg_folder, f"{base_name}.jpg")
            
            # 构建目标文件路径
            target_anno_file = os.path.join(target_anno_folder, f"{base_name}.xml")
            target_jpeg_file = os.path.join(target_jpeg_folder, f"{base_name}.jpg")
            
            # 复制文件
            if os.path.exists(source_anno_file):
                shutil.copy2(source_anno_file, target_anno_file)
                print(f"已复制标注文件: {base_name}.xml")
            else:
                print(f"警告: 未找到标注文件 {base_name}.xml")
                
            if os.path.exists(source_jpeg_file):
                shutil.copy2(source_jpeg_file, target_jpeg_file)
                print(f"已复制图片文件: {base_name}.jpg")
            else:
                print(f"警告: 未找到图片文件 {base_name}.jpg")

# 使用示例
if __name__ == "__main__":
    # 设置相关路径
    mask_folder = "/home/shuai/Documents/mydataset/my_usv_aug/EMBEDDED/mask"  # mask文件夹路径
    source_folder = "/home/shuai/Documents/mydataset/my_usv_aug/EMBEDDED_backup"  # 包含Annotations和JPEGImages的源文件夹路径
    target_base_folder = "/home/shuai/Documents/mydataset/my_usv_aug/EMBEDDED"  # 目标文件夹根目录路径
    
    # 处理文件
    process_files(mask_folder, source_folder, target_base_folder)